#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

void translate_main(CodeGen *codegen); // 将 main 函数翻译为汇编代码

int main() {
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 告诉汇编器将汇编放到代码段
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    translate_main(codegen);

    std::cout << codegen->print();
    delete codegen;
    delete module;
    return 0;
}

// TODO: 翻译你优化后的IR为RISC-V目标代码
void translate_main(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* 声明 main 函数 */ 
    codegen->append_inst(".align 2", ASMInstruction::Atrribute);
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);

    /* main 函数开始 */
    codegen->append_inst("main", ASMInstruction::Label);

    /* main 函数的 Prologue (序言) */
    codegen->append_inst("addi", {"sp", "sp", "-48"});
    codegen->append_inst("sd", {"ra", "sp", "40"});
    codegen->append_inst("sd", {"s0", "sp", "32"});
    codegen->append_inst("addi", {"s0", "sp", "48"});

    /* main 函数的 label_entry */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);
    codegen->append_inst("# br label %label1", ASMInstruction::Comment);
    codegen->append_inst("li", {"t0", "1"});  // 初始化phi节点的值为1
    codegen->append_inst("j", {".main_label1"});

    /* label1: phi node and loop condition */
    codegen->append_inst(".main_label1", ASMInstruction::Label);
    codegen->append_inst("# %op2 = phi i32 [ 1, %label_entry ], [ %op6, %label4 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op3 = icmp slt i32 %op2, 999999999", ASMInstruction::Comment);
    
    // 比较操作：检查t0是否小于999999999
    codegen->append_inst("li", {"t1", "999999999"});
    codegen->append_inst("slt", {"t2", "t0", "t1"});
    codegen->append_inst("# br i1 %op3, label %label4, label %label7", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t3", "zero", "0"});
    codegen->append_inst("blt", {"t3", "t2", ".main_label4"});
    codegen->append_inst("j", {".main_label7"});

    /* label4: loop body */
    codegen->append_inst(".main_label4", ASMInstruction::Label);
    codegen->append_inst("# %op6 = add i32 %op2, 1", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t0", "t0", "1"});  // t0 = t0 + 1，这个值会在下次循环中作为phi的输入
    codegen->append_inst("# br label %label1", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label1"});

    /* label7: loop exit */
    codegen->append_inst(".main_label7", ASMInstruction::Label);
    codegen->append_inst("# ret i32 %op2", ASMInstruction::Comment);
    codegen->append_inst("mv", {"a0", "t0"});  // 将结果移动到返回值寄存器

    /* main 函数的 Epilogue (收尾) */
    codegen->append_inst(".main_exit", ASMInstruction::Label);
    codegen->append_inst("ld", {"ra", "sp", "40"});
    codegen->append_inst("ld", {"s0", "sp", "32"});
    codegen->append_inst("addi", {"sp", "sp", "48"});
    codegen->append_inst("ret");
}
